import requests
import json
import time

def check_ollama_service():
    """检查 Ollama 服务是否正常运行"""
    try:
        response = requests.get("http://localhost:11434/api/tags")
        if response.status_code == 200:
            print("✓ Ollama 服务正常运行")
            return True
        else:
            print(f"✗ Ollama 服务响应异常: {response.status_code}")
            return False
    except Exception as e:
        print(f"✗ Ollama 服务检查失败: {str(e)}")
        return False

def generate_text(prompt, model="mistral", temperature=0.7, max_tokens=2000):
    """使用 Mistral 模型生成文本"""
    # Ollama API endpoint
    url = "http://localhost:11434/api/generate"
    
    # 准备请求数据
    data = {
        "model": model,
        "prompt": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": False
    }
    
    try:
        print(f"\n正在使用 {model} 模型生成文本...")
        print(f"提示词: {prompt}")
        
        # 发送请求
        response = requests.post(url, json=data)
        
        if response.status_code == 200:
            result = response.json()
            return result.get("response", "")
        else:
            print(f"请求失败: HTTP {response.status_code}")
            return None
            
    except Exception as e:
        print(f"发生错误: {str(e)}")
        return None

def main():
    print("=== Mistral 模型演示 ===")
    
    # 检查 Ollama 服务
    if not check_ollama_service():
        print("\n请先运行 'ollama serve' 启动服务")
        return
    
    # 示例提示词
    prompts = [
        {
            "prompt": "请用中文解释什么是人工智能，并给出三个实际应用例子。",
            "description": "基础问答"
        },
        {
            "prompt": "写一个Python函数，实现快速排序算法，并添加详细的中文注释。",
            "description": "代码生成"
        },
        {
            "prompt": "请总结以下文本的主要观点：人工智能正在改变我们的生活方式。从智能手机助手到自动驾驶汽车，AI技术已经渗透到我们日常生活的方方面面。在医疗领域，AI帮助医生更准确地诊断疾病；在教育领域，AI提供个性化的学习体验；在商业领域，AI优化运营效率。",
            "description": "文本摘要"
        }
    ]
    
    try:
        for i, prompt_data in enumerate(prompts, 1):
            print(f"\n=== 示例 {i}: {prompt_data['description']} ===")
            response = generate_text(prompt_data["prompt"])
            
            if response:
                print("\n模型响应:")
                print("-" * 50)
                print(response)
                print("-" * 50)
            else:
                print(f"示例 {i} 生成失败")
            
            # 添加延迟，避免请求过快
            if i < len(prompts):
                time.sleep(1)
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 